package com.loklok.code.confuse.util;

import cn.hutool.http.HttpUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;

import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import com.alibaba.fastjson.TypeReference;
import org.apache.commons.lang3.StringUtils;

import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;

public class SyncOpenAIUtil {

    private static final String regex = "```([\\s\\S]*?)```";
    private static final Pattern pattern = Pattern.compile(regex);

    private static final String WEN_XIN_YI_YAN_TOKEN_URI = "https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id=%s&client_secret=%s";
    private static final String WEN_XIN_YI_YAN_CHAT_URI = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=%s";

    public static String getWenXinYiYanToken(String clientId, String clientSecret){
        String uri = String.format(WEN_XIN_YI_YAN_TOKEN_URI, clientId, clientSecret);
        String res = HttpUtil.createPost(uri)
                .timeout(60000)
                .execute().body();
        if(StringUtils.isEmpty(res)){
            throw new RuntimeException("获取文心一言Token是空的");
        }
        JSONObject jsonObject = JSON.parseObject(res);
        return jsonObject.getString("access_token");
    }


    /**
     *  {
     *     "temperature": 1.0,
     *     "top_p": 0.8,
     *     "response_format":"json_object",
     *     "penalty_score": 1,
     *     "enable_system_memory": false,
     *     "disable_search": false,
     *     "enable_citation": false,
     *     "messages": [
     *         {
     *             "role":"user",
     *             "content": "请给出 [appBundleID、keyIdentifier、productIdentifier] 这些单词的英文同义词，要求每个单词的同义词返回1个，注意返回的结果不包括 [Application Bundle Identifier,Bundle Identifier] 这些词汇，结果案例如下形式: {'appBundleID':'appBundleID1', 'keyIdentifier':'keyIdentifier1'},只返回压缩的json，请注意不需要其他解释"
     *         }
     *     ]
     * }
     * @param accessToken
     */
    public static Map<String,String> getWenXinYiYanSynonym(String accessToken, String wenXinYiYanPrompt){
        JSONObject params = new JSONObject();
        params.put("temperature", 1.0);
        params.put("top_p", 0.8);
        params.put("response_format", "json_object");
        params.put("penalty_score", 1);
        params.put("enable_system_memory", false);
        params.put("disable_search", false);
        params.put("enable_citation", false);

        JSONArray messages = new JSONArray();
        JSONObject item = new JSONObject();
        item.put("role", "user");
        item.put("content", wenXinYiYanPrompt);
        messages.add(item);

        params.put("messages", messages);

        String uri = String.format(WEN_XIN_YI_YAN_CHAT_URI, accessToken);
        String res = HttpUtil.createPost(uri).
                header("Content-Type", "application/json")
                .body(params.toString().getBytes(StandardCharsets.UTF_8))
                .timeout(60000)
                .execute().body();

        if(StringUtils.isEmpty(res)){
            throw new RuntimeException("获取文心一言近义词是空的");
        }
        JSONObject jsonObject = JSON.parseObject(res);
        String result = jsonObject.getString("result" );
        result = result.replace("\\n", "").replace("\\", "");
        Map<String,String> map = JSON.parseObject(result, Map.class);
        return map;

    }

    public static Map<String,Set<String>> getSynonym(String accessToken, String prompt){
        String content = doDeepSeekRetryRequest(accessToken, prompt);
        // 正则表达式匹配 ``` 包裹的代码块
        Matcher matcher = pattern.matcher(content);
        // 查找并输出匹配的代码块
        if (matcher.find()) {
            String codeBlock = matcher.group(1); // 获取第一个捕获组（即代码内容）
            content = codeBlock.replace("json\n", "");
        }
        Map<String, Set<String>> map = JSON.parseObject(content,  new TypeReference<Map<String, Set<String>>>() {}.getType());
        return map;

    }

    public static String getVmBody(String accessToken, String prompt){
        String content = doDeepSeekRetryRequest(accessToken, prompt);
        Matcher matcher = pattern.matcher(content);
        // 查找并输出匹配的代码块
        if (matcher.find()) {
            String codeBlock = matcher.group(1); // 获取第一个捕获组（即代码内容）
            content = codeBlock.replace("vm\n", "");
        }
        return content;
    }

    public static String getParamBody(String accessToken, String prompt){
        String content = doDeepSeekRetryRequest(accessToken, prompt);
        // 正则表达式匹配 ``` 包裹的代码块
        Matcher matcher = pattern.matcher(content);
        // 查找并输出匹配的代码块
        if (matcher.find()) {
            String codeBlock = matcher.group(1); // 获取第一个捕获组（即代码内容）
            return codeBlock.replace("java\n", "");
        }
        return content;
    }

    private static String doDeepSeekRetryRequest(String apiKey, String content){
        int retry = 3;
        while (retry-- > 0) {
            try {
                return doDeepSeekRequest(apiKey, content);
            } catch (Exception e) {
                if (!e.getMessage().contains("timed out")) {
                    throw e;
                }
            }
        }
        throw new RuntimeException("请求 deepseek 超时多次失败");
    }

    private static String doDeepSeekRequest(String apiKey, String content){

        JSONObject jsonObject = new JSONObject();
        jsonObject.put("model", "deepseek-chat");
        jsonObject.put("stream", false);

        JSONArray jsonArray = new JSONArray();
        JSONObject messageObject = new JSONObject();
        messageObject.put("role", "user");
        messageObject.put("content", content.toString());
        jsonArray.add(messageObject);
        jsonObject.put("messages", jsonArray);
        Map<String, String> header = new HashMap<>();
        header.put("Authorization", "Bearer " + apiKey);
        header.put("Content-Type", "application/json");

        String res = HttpUtil.createPost("https://api.deepseek.com/chat/completions").
                headerMap(header, true)
                .body(JSONObject.toJSONString(jsonObject).getBytes(StandardCharsets.UTF_8))
                .timeout(1200000)
                .execute().body();

        JSONObject resJSONObject = JSONObject.parseObject(res);
        if(resJSONObject.containsKey("error_msg")){
            throw new RuntimeException("deepSeek error:" + resJSONObject.getString("error_msg"));
        }
        JSONObject choicesOne = JSONObject.class.cast(resJSONObject.getJSONArray("choices").get(0));

        if(!"stop".equals(choicesOne.getString("finish_reason"))){
            throw new RuntimeException("deep seek 回答未结束");
        }

        JSONObject message = choicesOne.getJSONObject("message");
        String resContent = message.getString("content");
        return resContent;
    }

}
